package parser;

import entity.capec.AttackPattern;
import entity.capec.ExecutionFlow;
import org.dom4j.Document;
import org.dom4j.Element;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import util.XmlFileUtil;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/**
 * 解析capec的工具类
 *
 * @author xym
 * @date 2020/7/15
 */
public class CapecXmlParser {
    private static final Logger log = LoggerFactory.getLogger(CapecXmlParser.class);

    /**
     * 对外暴露的解析接口，接收源xml文件输入，返回json文件中包含的攻击模式信息
     *
     * @param fileName xml源文件
     */
    public List<AttackPattern> execute(String fileName) {
        Document doc = XmlFileUtil.read(fileName);
        // 获取标签<Attack_Pattern_Catalog>
        Element root = doc.getRootElement();
        // 获取标签<Attack_Patterns>
        Element attackPatterns = root.element("Attack_Patterns");

        List<AttackPattern> result = new ArrayList<>();
        // 获取所有的<Attack_Pattern>标签，每一个标签对应于一个AttackPattern实体
        Iterator<Element> apIterator = attackPatterns.elementIterator("Attack_Pattern");
        while (apIterator.hasNext()) {
            Element singleAttackPattern = apIterator.next();
            // 调用函数，提取单个<Attack_Pattern>下的信息
            AttackPattern attackPattern = innerAttackPatternParsing(singleAttackPattern);
            if (attackPattern != null) {
                result.add(attackPattern);
            }
        }

        return result;
    }

    /**
     * 解析单个<Attack_Pattern>
     *
     * @param root <Attack_Pattern>
     */
    private AttackPattern innerAttackPatternParsing(Element root) {
        if (root == null) {
            return null;
        }

        AttackPattern apEntity = new AttackPattern();
        apEntity.setId(Integer.parseInt(root.attributeValue("ID")));
        apEntity.setName(root.attributeValue("Name"));

        Iterator<Element> subNodes = root.elementIterator();
        while (subNodes.hasNext()) {
            Element node = subNodes.next();
            String nodeName = node.getName();
            switch (nodeName) {
                case "Description":
                    apEntity.setDescription(node.getText());
                    break;
                case "Likelihood_Of_Attack":
                    apEntity.setLikelihoodOfAttack(node.getText());
                    break;
                case "Typical_Severity":
                    apEntity.setTypicalSeverity(node.getText());
                    break;
                case "Related_Attack_Patterns":
                    Iterator<Element> relatedAttackPatternIterator = node.elementIterator("Related_Attack_Pattern");
                    while (relatedAttackPatternIterator.hasNext()) {
                        Element relatedAttackPattern = relatedAttackPatternIterator.next();
                        String value = relatedAttackPattern.attributeValue("Nature")
                                + "-"
                                + relatedAttackPattern.attributeValue("CAPEC_ID");
                        apEntity.addRelatedAttackPattern(value);
                    }
                    break;
                case "Execution_Flow":
                    Iterator<Element> attackStepIterator = node.elementIterator("Attack_Step");
                    while (attackStepIterator.hasNext()) {
                        ExecutionFlow executionFlow = new ExecutionFlow();
                        Element attackStep = attackStepIterator.next();
                        Iterator<Element> subElementIterator = attackStep.elementIterator();
                        while (subElementIterator.hasNext()) {
                            Element subElement = subElementIterator.next();
                            switch (subElement.getName()) {
                                case "Step":
                                    executionFlow.setStep(Integer.parseInt(subElement.getText()));
                                    break;
                                case "Phase":
                                    executionFlow.setPhaseName(subElement.getText());
                                    break;
                                case "Description":
                                    executionFlow.setDescription(subElement.getText());
                                    break;
                                case "Technique":
                                    executionFlow.addTechniques(subElement.getText());
                                    break;
                                default:
                                    break;
                            }
                        }

                        apEntity.addExecutionFlow(executionFlow);
                    }
                    break;
                case "Prerequisites":
                    Iterator<Element> prerequisiteIterator = node.elementIterator("Prerequisite");
                    while (prerequisiteIterator.hasNext()) {
                        Element subElement = prerequisiteIterator.next();
                        apEntity.addPrerequisites(subElement.getText());
                    }
                    break;
                case "Skills_Required":
                    Iterator<Element> skillsRequiredIterator = node.elementIterator("Skill");
                    while (skillsRequiredIterator.hasNext()) {
                        Element subElement = skillsRequiredIterator.next();
                        apEntity.addSkillsRequired(subElement.getText());
                    }
                    break;
                case "Resources_Required":
                    Iterator<Element> resourceIterator = node.elementIterator("Resource");
                    while (resourceIterator.hasNext()) {
                        Element subElement = resourceIterator.next();
                        apEntity.addResourcesRequired(subElement.getText());
                    }
                    break;
                case "Indicators":
                    Iterator<Element> indicatorIterator = node.elementIterator("Indicator");
                    while (indicatorIterator.hasNext()) {
                        Element subElement = indicatorIterator.next();
                        apEntity.addIndicator(subElement.getText());
                    }
                    break;
                case "Consequences":
                    Iterator<Element> consequenceIterator = node.elementIterator("Consequence");
                    while (consequenceIterator.hasNext()) {
                        Element subElement = consequenceIterator.next();
                        Element impact = subElement.element("Impact");
                        if (impact != null) {
                            apEntity.addConsequence(impact.getText());
                        }
                    }
                    break;
                case "Mitigations":
                    Iterator<Element> mitigationIterator = node.elementIterator("Mitigation");
                    while (mitigationIterator.hasNext()) {
                        Element subElement = mitigationIterator.next();
                        if (subElement.elementIterator().hasNext()) {
                            Iterator<Element> ssubElementIterator = subElement.elementIterator();
                            while (ssubElementIterator.hasNext()) {
                                apEntity.addMitigation(ssubElementIterator.next().getText());
                            }
                        } else {
                            apEntity.addMitigation(subElement.getText());
                        }
                    }
                    break;
                case "Related_Weaknesses":
                    Iterator<Element> relatedWeaknessIterator = node.elementIterator("Related_Weakness");
                    while (relatedWeaknessIterator.hasNext()) {
                        Element subElement = relatedWeaknessIterator.next();
                        apEntity.addRelatedWeakness(Integer.parseInt(subElement.attributeValue("CWE_ID")));
                    }
                    break;
                default:
                    break;
            }
        }

        return apEntity;
    }
}
